import javax.crypto.Cipher;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import javax.swing.*;
import java.awt.*;
import java.awt.event.KeyEvent;
import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.security.SecureRandom;

public class AESFileEncryptorGUI {
    // AES算法名称和工作模式
    private static final String AES_ALGORITHM = "AES/CBC/PKCS5Padding";
    // 密钥长度（256位）
    private static final int KEY_SIZE = 256;
    // IV长度（16字节）
    private static final int IV_SIZE = 16;

    // ANSI 颜色代码(用于控制台)
    private static final String RESET = "\u001B[0m";
    private static final String BLACK = "\u001B[30m";
    private static final String RED = "\u001B[31m";
    private static final String GREEN = "\u001B[32m";
    private static final String YELLOW = "\u001B[33m";
    private static final String BLUE = "\u001B[34m";
    private static final String PURPLE = "\u001B[35m";
    private static final String CYAN = "\u001B[36m";
    private static final String WHITE = "\u001B[37m";

    public static void main(String[] args) {
        // 创建并显示GUI
        SwingUtilities.invokeLater(() -> {
            AESApplicationFrame frame = new AESApplicationFrame();
            frame.setVisible(true);
        });
    }

    // ========== GUI 类定义 ==========

    /**
     * 主应用程序窗口
     */
    static class AESApplicationFrame extends JFrame {
        public AESApplicationFrame() {
            initUI();
        }

        private void initUI() {
            setTitle("AES-256 文件加密系统");
            setSize(800, 600);
            setDefaultCloseOperation(EXIT_ON_CLOSE);
            setLocationRelativeTo(null); // 居中显示

            // 创建主面板
            JPanel mainPanel = new JPanel(new BorderLayout());

            // 创建顶部标题面板
            JPanel titlePanel = new JPanel();
            JLabel titleLabel = new JLabel("欢迎使用 AES-256 文件加密系统");
            titleLabel.setFont(new Font("Arial", Font.BOLD, 24));
            titleLabel.setForeground(Color.BLUE);
            titlePanel.add(titleLabel);

            // 创建菜单栏
            JMenuBar menuBar = new JMenuBar();

            // 主菜单
            JMenu mainMenu = new JMenu("操作");
            mainMenu.setMnemonic(KeyEvent.VK_O);

            // 生成密钥菜单项
            JMenuItem generateKeyItem = new JMenuItem("生成密钥", KeyEvent.VK_G);
            generateKeyItem.addActionListener(e -> showGenerateKeyDialog());

            // 加密文件菜单项
            JMenuItem encryptItem = new JMenuItem("加密文件", KeyEvent.VK_E);
            encryptItem.addActionListener(e -> showEncryptDialog());

            // 解密文件菜单项
            JMenuItem decryptItem = new JMenuItem("解密文件", KeyEvent.VK_D);
            decryptItem.addActionListener(e -> showDecryptDialog());

            // 退出菜单项
            JMenuItem exitItem = new JMenuItem("退出", KeyEvent.VK_X);
            exitItem.addActionListener(e -> System.exit(0));

            // 添加菜单项到主菜单
            mainMenu.add(generateKeyItem);
            mainMenu.add(encryptItem);
            mainMenu.add(decryptItem);
            mainMenu.addSeparator();
            mainMenu.add(exitItem);

            // 添加菜单栏
            menuBar.add(mainMenu);
            setJMenuBar(menuBar);

            // 创建内容面板
            JPanel contentPanel = new JPanel(new GridBagLayout());
            GridBagConstraints gbc = new GridBagConstraints();
            gbc.gridwidth = GridBagConstraints.REMAINDER;
            gbc.anchor = GridBagConstraints.NORTH;
            gbc.insets = new Insets(10, 10, 10, 10);

            // 添加欢迎标签
            JTextArea welcomeText = new JTextArea("欢迎使用 AES-256 文件加密系统\n\n" +
                    "本软件提供以下功能：\n" +
                    "1. 生成并保存 AES-256 密钥\n" +
                    "2. 使用 AES-256 CBC 模式加密文件\n" +
                    "3. 使用 AES-256 CBC 模式解密文件\n\n" +
                    "请通过上方菜单选择操作");
            welcomeText.setEditable(false);
            welcomeText.setOpaque(false);
            welcomeText.setFont(new Font("Arial", Font.PLAIN, 16));
            welcomeText.setPreferredSize(new Dimension(600, 200));
            contentPanel.add(welcomeText, gbc);

            // 添加操作按钮
            JPanel buttonPanel = new JPanel(new FlowLayout(FlowLayout.CENTER, 20, 10));

            JButton generateKeyButton = new JButton("生成密钥");
            generateKeyButton.addActionListener(e -> showGenerateKeyDialog());
            generateKeyButton.setPreferredSize(new Dimension(150, 40));

            JButton encryptButton = new JButton("加密文件");
            encryptButton.addActionListener(e -> showEncryptDialog());
            encryptButton.setPreferredSize(new Dimension(150, 40));

            JButton decryptButton = new JButton("解密文件");
            decryptButton.addActionListener(e -> showDecryptDialog());
            decryptButton.setPreferredSize(new Dimension(150, 40));

            buttonPanel.add(generateKeyButton);
            buttonPanel.add(encryptButton);
            buttonPanel.add(decryptButton);

            contentPanel.add(buttonPanel, gbc);

            // 添加组件到主面板
            mainPanel.add(titlePanel, BorderLayout.NORTH);
            mainPanel.add(contentPanel, BorderLayout.CENTER);

            add(mainPanel);
        }

        private void showGenerateKeyDialog() {
            GenerateKeyDialog dialog = new GenerateKeyDialog(this);
            dialog.setVisible(true);
        }

        private void showEncryptDialog() {
            EncryptDialog dialog = new EncryptDialog(this);
            dialog.setVisible(true);
        }

        private void showDecryptDialog() {
            DecryptDialog dialog = new DecryptDialog(this);
            dialog.setVisible(true);
        }
    }

    /**
     * 生成密钥对话框
     */
    static class GenerateKeyDialog extends JDialog {
        private JTextField filePathField;
        private JPasswordField passwordField;

        public GenerateKeyDialog(JFrame parent) {
            super(parent, "生成密钥", true);
            initUI();
        }

        private void initUI() {
            setSize(500, 200);
            setLocationRelativeTo(getParent());

            JPanel panel = new JPanel(new GridLayout(4, 2, 10, 10));

            panel.add(new JLabel("保存路径:"));
            filePathField = new JTextField();
            panel.add(filePathField);

            JButton choosePathButton = new JButton("选择路径");
            choosePathButton.addActionListener(e -> chooseFilePath());
            panel.add(new JLabel()); // 空单元格
            panel.add(choosePathButton);

            panel.add(new JLabel("密码 (可选):"));
            passwordField = new JPasswordField();
            panel.add(passwordField);

            JButton generateButton = new JButton("生成密钥");
            generateButton.addActionListener(e -> generateKey());
            JButton cancelButton = new JButton("取消");
            cancelButton.addActionListener(e -> dispose());

            JPanel buttonPanel = new JPanel(new FlowLayout(FlowLayout.RIGHT));
            buttonPanel.add(generateButton);
            buttonPanel.add(cancelButton);

            add(panel, BorderLayout.CENTER);
            add(buttonPanel, BorderLayout.SOUTH);
        }

        private void chooseFilePath() {
            JFileChooser fileChooser = new JFileChooser();
            fileChooser.setDialogTitle("选择密钥保存位置");
            fileChooser.setFileSelectionMode(JFileChooser.FILES_ONLY);

            if (fileChooser.showSaveDialog(this) == JFileChooser.APPROVE_OPTION) {
                filePathField.setText(fileChooser.getSelectedFile().getAbsolutePath());
            }
        }

        private void generateKey() {
            String keyFilePath = filePathField.getText().trim();

            if (keyFilePath.isEmpty()) {
                JOptionPane.showMessageDialog(this, "错误: 密钥文件路径不能为空！",
                        "错误", JOptionPane.ERROR_MESSAGE);
                return;
            }

            File keyFile = new File(keyFilePath);
            if (keyFile.exists()) {
                int response = JOptionPane.showConfirmDialog(this,
                        "警告: 密钥文件已存在。是否覆盖?", "确认覆盖",
                        JOptionPane.YES_NO_OPTION, JOptionPane.WARNING_MESSAGE);

                if (response != JOptionPane.YES_OPTION) {
                    return;
                }
            }

            try {
                // 生成随机密钥
                byte[] key = generateKeyBYTE();

                // 如果有密码，使用密码加密存储(这里暂未实现加密存储功能)
                // 保存密钥到文件
                saveKeyToFile(key, keyFilePath);

                JOptionPane.showMessageDialog(this,
                        "成功: 密钥已生成并保存到 " + keyFilePath,
                        "成功", JOptionPane.INFORMATION_MESSAGE);

                dispose();
            } catch (Exception e) {
                JOptionPane.showMessageDialog(this,
                        "发生错误: " + e.getMessage(),
                        "错误", JOptionPane.ERROR_MESSAGE);
                e.printStackTrace();
            }
        }

        private byte[] generateKeyBYTE() throws Exception {
            SecureRandom random = new SecureRandom();
            byte[] key = new byte[KEY_SIZE / 8]; // 256位=32字节
            random.nextBytes(key);
            return key;
        }

        private void saveKeyToFile(byte[] key, String filePath) throws IOException {
            try (FileOutputStream fos = new FileOutputStream(filePath)) {
                fos.write(key);
            }
        }
    }

    /**
     * 加密文件对话框
     */
    static class EncryptDialog extends JDialog {
        private JTextField keyPathField;
        private JTextField inputPathField;
        private JTextField outputPathField;

        public EncryptDialog(JFrame parent) {
            super(parent, "文件加密", true);
            initUI();
        }

        private void initUI() {
            setSize(600, 300);
            setLocationRelativeTo(getParent());

            JPanel panel = new JPanel(new GridLayout(5, 2, 10, 10));

            panel.add(new JLabel("密钥文件:"));
            keyPathField = new JTextField();
            panel.add(keyPathField);
            JButton chooseKeyButton = new JButton("选择文件");
            chooseKeyButton.addActionListener(e -> chooseKeyPath());
            panel.add(new JLabel());
            panel.add(chooseKeyButton);

            panel.add(new JLabel("要加密的文件:"));
            inputPathField = new JTextField();
            panel.add(inputPathField);
            JButton chooseInputButton = new JButton("选择文件");
            chooseInputButton.addActionListener(e -> chooseInputPath());
            panel.add(new JLabel());
            panel.add(chooseInputButton);

            panel.add(new JLabel("保存加密结果至:"));
            outputPathField = new JTextField();
            panel.add(outputPathField);
            JButton chooseOutputButton = new JButton("选择文件");
            chooseOutputButton.addActionListener(e -> chooseOutputPath());
            panel.add(new JLabel());
            panel.add(chooseOutputButton);

            JButton encryptButton = new JButton("开始加密");
            encryptButton.addActionListener(e -> encrypt());
            JButton cancelButton = new JButton("取消");
            cancelButton.addActionListener(e -> dispose());

            JPanel buttonPanel = new JPanel(new FlowLayout(FlowLayout.RIGHT));
            buttonPanel.add(encryptButton);
            buttonPanel.add(cancelButton);

            add(panel, BorderLayout.CENTER);
            add(buttonPanel, BorderLayout.SOUTH);
        }

        private void chooseKeyPath() {
            JFileChooser fileChooser = new JFileChooser();
            fileChooser.setDialogTitle("选择密钥文件");
            fileChooser.setFileSelectionMode(JFileChooser.FILES_ONLY);

            if (fileChooser.showOpenDialog(this) == JFileChooser.APPROVE_OPTION) {
                keyPathField.setText(fileChooser.getSelectedFile().getAbsolutePath());
            }
        }

        private void chooseInputPath() {
            JFileChooser fileChooser = new JFileChooser();
            fileChooser.setDialogTitle("选择要加密的文件");
            fileChooser.setFileSelectionMode(JFileChooser.FILES_ONLY);

            if (fileChooser.showOpenDialog(this) == JFileChooser.APPROVE_OPTION) {
                inputPathField.setText(fileChooser.getSelectedFile().getAbsolutePath());
            }
        }

        private void chooseOutputPath() {
            JFileChooser fileChooser = new JFileChooser();
            fileChooser.setDialogTitle("选择加密文件保存位置");
            fileChooser.setFileSelectionMode(JFileChooser.FILES_ONLY);

            if (fileChooser.showSaveDialog(this) == JFileChooser.APPROVE_OPTION) {
                outputPathField.setText(fileChooser.getSelectedFile().getAbsolutePath());
            }
        }

        private void encrypt() {
            String keyFilePath = keyPathField.getText().trim();
            String inputFile = inputPathField.getText().trim();
            String outputFile = outputPathField.getText().trim();

            if (keyFilePath.isEmpty() || inputFile.isEmpty() || outputFile.isEmpty()) {
                JOptionPane.showMessageDialog(this,
                        "错误: 所有路径都不能为空！",
                        "错误", JOptionPane.ERROR_MESSAGE);
                return;
            }

            try {
                // 加载密钥
                byte[] key = loadKeyFromFile(keyFilePath);

                // 加密文件
                encryptFile(key, inputFile, outputFile);

                JOptionPane.showMessageDialog(this,
                        "成功: 文件加密完成。加密结果已保存到 " + outputFile,
                        "成功", JOptionPane.INFORMATION_MESSAGE);

                dispose();
            } catch (Exception e) {
                JOptionPane.showMessageDialog(this,
                        "发生错误: " + e.getMessage(),
                        "错误", JOptionPane.ERROR_MESSAGE);
                e.printStackTrace();
            }
        }

        private byte[] loadKeyFromFile(String filePath) throws IOException {
            return Files.readAllBytes(Paths.get(filePath));
        }

        private void encryptFile(byte[] key, String inputFile, String outputFile) throws Exception {
            // 读取原始文件内容
            byte[] fileContent = Files.readAllBytes(Paths.get(inputFile));

            // 创建CBC模式使用的IV参数
            byte[] iv = new byte[IV_SIZE];
            new SecureRandom().nextBytes(iv);

            // 创建AES密钥规格
            SecretKeySpec keySpec = new SecretKeySpec(key, "AES");
            IvParameterSpec ivSpec = new IvParameterSpec(iv);

            // 创建并初始化加密器
            Cipher cipher = Cipher.getInstance(AES_ALGORITHM);
            cipher.init(Cipher.ENCRYPT_MODE, keySpec, ivSpec);

            // 执行加密
            byte[] encrypted = cipher.doFinal(fileContent);

            // 将IV和加密后的数据合并，写入输出文件
            try (FileOutputStream fos = new FileOutputStream(outputFile)) {
                fos.write(iv); // 先写入IV
                fos.write(encrypted); // 再写入加密数据
            }
        }
    }

    /**
     * 解密文件对话框
     */
    static class DecryptDialog extends JDialog {
        private JTextField keyPathField;
        private JTextField inputPathField;
        private JTextField outputPathField;

        public DecryptDialog(JFrame parent) {
            super(parent, "文件解密", true);
            initUI();
        }

        private void initUI() {
            setSize(600, 300);
            setLocationRelativeTo(getParent());

            JPanel panel = new JPanel(new GridLayout(5, 2, 10, 10));

            panel.add(new JLabel("密钥文件:"));
            keyPathField = new JTextField();
            panel.add(keyPathField);
            JButton chooseKeyButton = new JButton("选择文件");
            chooseKeyButton.addActionListener(e -> chooseKeyPath());
            panel.add(new JLabel());
            panel.add(chooseKeyButton);

            panel.add(new JLabel("要解密的文件:"));
            inputPathField = new JTextField();
            panel.add(inputPathField);
            JButton chooseInputButton = new JButton("选择文件");
            chooseInputButton.addActionListener(e -> chooseInputPath());
            panel.add(new JLabel());
            panel.add(chooseInputButton);

            panel.add(new JLabel("保存解密结果至:"));
            outputPathField = new JTextField();
            panel.add(outputPathField);
            JButton chooseOutputButton = new JButton("选择文件");
            chooseOutputButton.addActionListener(e -> chooseOutputPath());
            panel.add(new JLabel());
            panel.add(chooseOutputButton);

            JButton decryptButton = new JButton("开始解密");
            decryptButton.addActionListener(e -> decrypt());
            JButton cancelButton = new JButton("取消");
            cancelButton.addActionListener(e -> dispose());

            JPanel buttonPanel = new JPanel(new FlowLayout(FlowLayout.RIGHT));
            buttonPanel.add(decryptButton);
            buttonPanel.add(cancelButton);

            add(panel, BorderLayout.CENTER);
            add(buttonPanel, BorderLayout.SOUTH);
        }

        private void chooseKeyPath() {
            JFileChooser fileChooser = new JFileChooser();
            fileChooser.setDialogTitle("选择密钥文件");
            fileChooser.setFileSelectionMode(JFileChooser.FILES_ONLY);

            if (fileChooser.showOpenDialog(this) == JFileChooser.APPROVE_OPTION) {
                keyPathField.setText(fileChooser.getSelectedFile().getAbsolutePath());
            }
        }

        private void chooseInputPath() {
            JFileChooser fileChooser = new JFileChooser();
            fileChooser.setDialogTitle("选择要解密的文件");
            fileChooser.setFileSelectionMode(JFileChooser.FILES_ONLY);

            if (fileChooser.showOpenDialog(this) == JFileChooser.APPROVE_OPTION) {
                inputPathField.setText(fileChooser.getSelectedFile().getAbsolutePath());
            }
        }

        private void chooseOutputPath() {
            JFileChooser fileChooser = new JFileChooser();
            fileChooser.setDialogTitle("选择解密文件保存位置");
            fileChooser.setFileSelectionMode(JFileChooser.FILES_ONLY);

            if (fileChooser.showSaveDialog(this) == JFileChooser.APPROVE_OPTION) {
                outputPathField.setText(fileChooser.getSelectedFile().getAbsolutePath());
            }
        }

        private void decrypt() {
            String keyFilePath = keyPathField.getText().trim();
            String inputFile = inputPathField.getText().trim();
            String outputFile = outputPathField.getText().trim();

            if (keyFilePath.isEmpty() || inputFile.isEmpty() || outputFile.isEmpty()) {
                JOptionPane.showMessageDialog(this,
                        "错误: 所有路径都不能为空！",
                        "错误", JOptionPane.ERROR_MESSAGE);
                return;
            }

            try {
                // 加载密钥
                byte[] key = loadKeyFromFile(keyFilePath);

                // 解密文件
                decryptFile(key, inputFile, outputFile);

                JOptionPane.showMessageDialog(this,
                        "成功: 文件解密完成。解密结果已保存到 " + outputFile,
                        "成功", JOptionPane.INFORMATION_MESSAGE);

                dispose();
            } catch (Exception e) {
                JOptionPane.showMessageDialog(this,
                        "发生错误: " + e.getMessage(),
                        "错误", JOptionPane.ERROR_MESSAGE);
                e.printStackTrace();
            }
        }

        private byte[] loadKeyFromFile(String filePath) throws IOException {
            return Files.readAllBytes(Paths.get(filePath));
        }

        private void decryptFile(byte[] key, String inputFile, String outputFile) throws Exception {
            // 读取加密文件内容
            byte[] encryptedData = Files.readAllBytes(Paths.get(inputFile));

            // 提取IV（前16字节）
            byte[] iv = new byte[IV_SIZE];
            System.arraycopy(encryptedData, 0, iv, 0, IV_SIZE);
            IvParameterSpec ivSpec = new IvParameterSpec(iv);

            // 提取实际加密数据（剩余部分）
            int encryptedFileSize = encryptedData.length - IV_SIZE;
            byte[] encryptedContent = new byte[encryptedFileSize];
            System.arraycopy(encryptedData, IV_SIZE, encryptedContent, 0, encryptedFileSize);

            // 创建AES密钥规格
            SecretKeySpec keySpec = new SecretKeySpec(key, "AES");

            // 创建并初始化解密器
            Cipher cipher = Cipher.getInstance(AES_ALGORITHM);
            cipher.init(Cipher.DECRYPT_MODE, keySpec, ivSpec);

            // 执行解密
            byte[] decrypted = cipher.doFinal(encryptedContent);

            // 写入解密后的内容到输出文件
            try (FileOutputStream fos = new FileOutputStream(outputFile)) {
                fos.write(decrypted);
            }
        }
    }

    // 原始的静态方法保持不变
    /**
     * 生成256位的AES随机密钥
     */
    private static byte[] generateKey() throws Exception {
        SecureRandom random = new SecureRandom();
        byte[] key = new byte[KEY_SIZE / 8]; // 256位=32字节
        random.nextBytes(key);
        return key;
    }

    /**
     * 将密钥保存到文件
     */
    private static void saveKeyToFile(byte[] key, String filePath) throws IOException {
        try (FileOutputStream fos = new FileOutputStream(filePath)) {
            fos.write(key);
        }
    }

    /**
     * 从文件加载AES密钥
     */
    private static byte[] loadKeyFromFile(String filePath) throws IOException {
        return Files.readAllBytes(Paths.get(filePath));
    }

    /**
     * 使用AES加密文件
     */
    private static void encryptFile(byte[] key, String inputFile, String outputFile) throws Exception {
        // 读取原始文件内容
        byte[] fileContent = Files.readAllBytes(Paths.get(inputFile));

        // 创建CBC模式使用的IV参数
        byte[] iv = new byte[IV_SIZE];
        new SecureRandom().nextBytes(iv);

        // 创建AES密钥规格
        SecretKeySpec keySpec = new SecretKeySpec(key, "AES");
        IvParameterSpec ivSpec = new IvParameterSpec(iv);

        // 创建并初始化加密器
        Cipher cipher = Cipher.getInstance(AES_ALGORITHM);
        cipher.init(Cipher.ENCRYPT_MODE, keySpec, ivSpec);

        // 执行加密
        byte[] encrypted = cipher.doFinal(fileContent);

        // 将IV和加密后的数据合并，写入输出文件
        try (FileOutputStream fos = new FileOutputStream(outputFile)) {
            fos.write(iv); // 先写入IV
            fos.write(encrypted); // 再写入加密数据
        }
    }

    /**
     * 使用AES解密文件
     */
    private static void decryptFile(byte[] key, String inputFile, String outputFile) throws Exception {
        // 读取加密文件内容
        byte[] encryptedData = Files.readAllBytes(Paths.get(inputFile));

        // 提取IV（前16字节）
        byte[] iv = new byte[IV_SIZE];
        System.arraycopy(encryptedData, 0, iv, 0, IV_SIZE);
        IvParameterSpec ivSpec = new IvParameterSpec(iv);

        // 提取实际加密数据（剩余部分）
        int encryptedFileSize = encryptedData.length - IV_SIZE;
        byte[] encryptedContent = new byte[encryptedFileSize];
        System.arraycopy(encryptedData, IV_SIZE, encryptedContent, 0, encryptedFileSize);

        // 创建AES密钥规格
        SecretKeySpec keySpec = new SecretKeySpec(key, "AES");

        // 创建并初始化解密器
        Cipher cipher = Cipher.getInstance(AES_ALGORITHM);
        cipher.init(Cipher.DECRYPT_MODE, keySpec, ivSpec);

        // 执行解密
        byte[] decrypted = cipher.doFinal(encryptedContent);

        // 写入解密后的内容到输出文件
        try (FileOutputStream fos = new FileOutputStream(outputFile)) {
            fos.write(decrypted);
        }
    }
}